import os
import codecs
import numpy as np
import torch

__all__ = ['load_data']


SN3_PASCALVINCENT_TYPEMAP = {
    8: torch.uint8,
    9: torch.int8,
    11: torch.int16,
    12: torch.int32,
    13: torch.float32,
    14: torch.float64,
}


def read_feature_file(path: str) -> torch.Tensor:
    x = read_sn3_pascalvincent_tensor(path, strict=False)
    if x.dtype != torch.float32:
        raise TypeError(
            f"x should be of dtype torch.uint8 instead of {x.dtype}")
    if x.ndimension() != 4:
        raise ValueError(
            f"x should have 4 dimension instead of {x.ndimension()}")
    return x


def read_target_file(path: str) -> torch.Tensor:
    x = read_sn3_pascalvincent_tensor(path, strict=False)
    if x.dtype != torch.uint8:
        raise TypeError(
            f"x should be of dtype torch.uint8 instead of {x.dtype}")
    if x.ndimension() != 1:
        raise ValueError(
            f"x should have 1 dimension instead of {x.ndimension()}")
    return x.long()


def read_sn3_pascalvincent_tensor(path: str, strict: bool = True) -> torch.Tensor:
    """
    Read a SN3 file in "Pascal Vincent" format (Lush file 'libidx/idx-io.lsh').
    Argument may be a filename, compressed filename, or file object.
    """
    # read
    with open(path, "rb") as f:
        data = f.read()

    # parse
    magic = get_int(data[0:4])
    nd = magic % 256
    ty = magic // 256
    assert 1 <= nd <= 4
    assert 8 <= ty <= 14
    torch_type = SN3_PASCALVINCENT_TYPEMAP[ty]

    # nd = 4 (I, D, H, W) = (1200, 12, 9, 9)
    s = [get_int(data[4 * (i + 1): 4 * (i + 2)]) for i in range(nd)]
    # s[0] = 60000

    # The MNIST format uses the big endian byte order. If the system uses little endian byte order by default,
    # we need to reverse the bytes before we can read them with torch.frombuffer().
    # num_bytes_per_value = torch.iinfo(torch_type).bits // 8
    # num_bytes_per_value = torch.finfo(torch_type).bits // 8
    # needs_byte_reversal = sys.byteorder == "little" and num_bytes_per_value > 1

    parsed = torch.frombuffer(
        bytearray(data), dtype=torch_type, offset=(4 * (nd + 1)))
    # if needs_byte_reversal:
    #     parsed = parsed.flip(0)
    assert parsed.shape[0] == np.prod(s) or not strict
    return parsed.view(*s)


def get_int(b: bytes) -> int:
    return int(codecs.encode(b, "hex"), 16)


def load_data(feature_file: str, target_file: str) -> tuple[torch.Tensor, torch.Tensor]:
    """
    装载数据

    参数:
        feature_file:   特征数据文件全路径
        target_file:    目标(标签)文件全路径

    返回:
        特征数据, 目标数据
    """
    features = read_feature_file(feature_file)
    targets = read_target_file(target_file)
    return features, targets
